{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Formulate Instruct-Tuning Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate Task-Response with GPT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from Ashare_data import *\n",
    "import pandas as pd\n",
    "import akshare as ak\n",
    "import pandas as pd\n",
    "import os\n",
    "import csv\n",
    "import re\n",
    "import time\n",
    "import math\n",
    "import json\n",
    "import random\n",
    "from datasets import Dataset\n",
    "import datasets\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openai import OpenAI\n",
    "\n",
    "# client = OpenAI(api_key = \"YOURAPI\")\n",
    "\n",
    "start_date = \"20230201\"\n",
    "end_date = \"20240101\"\n",
    "DATA_DIR = f\"./{start_date}_{end_date}_24Apr_qfq\"\n",
    "os.makedirs(DATA_DIR, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def append_to_csv(filename, input_data, output_data):\n",
    "    \n",
    "    with open(filename, mode='a', newline='') as file:\n",
    "        writer = csv.writer(file)\n",
    "        writer.writerow([input_data, output_data])\n",
    "\n",
    "        \n",
    "def initialize_csv(filename):\n",
    "    \n",
    "    with open(filename, mode='w', newline='') as file:\n",
    "        writer = csv.writer(file)\n",
    "        writer.writerow([\"prompt\", \"answer\"])\n",
    "\n",
    "def query_gpt4(symbol_list, min_past_weeks=1, max_past_weeks=2, with_basics=True):\n",
    "\n",
    "    for symbol in symbol_list:\n",
    "        \n",
    "        csv_file = f'{DATA_DIR}/{symbol}_{start_date}_{end_date}_gpt-4.csv' if with_basics else \\\n",
    "                   f'{DATA_DIR}/{symbol}_{start_date}_{end_date}_nobasics_gpt-4.csv'\n",
    "        \n",
    "        if not os.path.exists(csv_file):\n",
    "            initialize_csv(csv_file)\n",
    "            pre_done = 0\n",
    "        else:\n",
    "            df = pd.read_csv(csv_file)\n",
    "            pre_done = len(df)\n",
    "\n",
    "        prompts = get_all_prompts_new(symbol, min_past_weeks, max_past_weeks, with_basics)\n",
    "\n",
    "        for i, prompt in enumerate(prompts):\n",
    "            \n",
    "            if i < pre_done:\n",
    "                continue\n",
    "\n",
    "            print(f\"{symbol} - {i}\")\n",
    "            \n",
    "            cnt = 0\n",
    "            while cnt < 5:\n",
    "                try:\n",
    "                    completion = client.chat.completions.create(\n",
    "                        model=\"gpt-4\",\n",
    "                        messages=[\n",
    "                            {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
    "                            {\"role\": \"user\", \"content\": prompt}\n",
    "                          ]\n",
    "                    )\n",
    "                    print(\"==Generate answer successfully==\")\n",
    "                    break    \n",
    "                except Exception:\n",
    "                    cnt += 1\n",
    "                    print(f'retry cnt {cnt}')\n",
    "            \n",
    "            answer = completion.choices[0].message.content if cnt < 5 else \"\"\n",
    "            append_to_csv(csv_file, prompt, answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "SZ50 = pd.read_excel(\"000016.SH-成分及权重-20240411.xlsx\")\n",
    "tickers = [tk[:6] for tk in SZ50['代码'].tolist()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "HS300 = pd.read_csv(\"HS300Index.csv\", header=None).iloc[:,0:2]\n",
    "HS300.columns = ['symbol', 'name']\n",
    "HS300.symbol = HS300.symbol.apply(lambda x: \"0\"*(6-len(str(x)))+str(x))\n",
    "HS_index = HS300.symbol.tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "688041\n"
     ]
    }
   ],
   "source": [
    "for i in tickers:\n",
    "    if i not in HS_index:\n",
    "        print(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "tickers.remove(\"688041\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_gpt4(tickers[:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_DIR = '20230201_20240101_24Apr_qfq'\n",
    "B_INST, E_INST = \"[INST]\", \"[/INST]\"\n",
    "B_SYS, E_SYS = \"<<SYS>>\\n\", \"\\n<</SYS>>\\n\\n\"\n",
    "\n",
    "def gpt4_to_llama_Chinese_new(symbol, with_basics=True):\n",
    "    \n",
    "    csv_file = f'{DATA_DIR}/{symbol}_{start_date}_{end_date}_gpt-4.csv' if with_basics else \\\n",
    "                   f'{DATA_DIR}/{symbol}_{start_date}_{end_date}_nobasics_gpt-4.csv'\n",
    "    \n",
    "    df = pd.read_csv(csv_file)\n",
    "    \n",
    "    prompts, answers, periods, labels = [], [], [], []\n",
    "    \n",
    "    for i, row in df.iterrows():\n",
    "        \n",
    "        prompt, answer = row['prompt'], row['answer']\n",
    "        \n",
    "        res = re.search(r\"那么让我们假设你对于下一周\\((.*)\\)的预测是((:?上涨|下跌).*%|股价持平)。\", prompt)\n",
    "        try:\n",
    "            period, label = res.group(1), res.group(2)\n",
    "        except AttributeError:\n",
    "            # set this to check for unchanged price(if the data formulation form has been updated, then this will be skipped)\n",
    "            res = re.search(r\"那么让我们假设你对于下一周\\((.*)\\)的预测是((:?上涨|下跌).)。\", prompt)\n",
    "            period, label = res.group(1), \"股价持平\"\n",
    "        \n",
    "        prompt = re.sub(\n",
    "                    r\"那么让我们假设你对于下一周\\((.*)\\)的预测是((:?上涨|下跌).*%|股价持平)。提供一个总结分析来支持你的预测。预测结果需要从你最后的分析中推断出来，因此不作为你分析的基础因素。\", \n",
    "                    f\"接下来请预测{symbol}下周({period})的股票涨跌幅，并提供一个总结分析来支持你的预测。\",\n",
    "                    prompt\n",
    "                )\n",
    "        try:\n",
    "            answer = re.sub(\n",
    "                r\"\\[预测和分析\\]：\\n\",\n",
    "                f\"[预测和分析]：\\n预测涨跌幅：{label}\\n总结分析：\",\n",
    "                answer\n",
    "            )\n",
    "        except Exception:\n",
    "            print(symbol, i)\n",
    "            print(label)\n",
    "            print(answer)\n",
    "            continue\n",
    "            \n",
    "        new_system_prompt = SYSTEM_PROMPT.replace('：\\n...', '：\\n预测涨跌幅：...\\n总结分析：...')\n",
    "        \n",
    "        prompt = B_INST + B_SYS + new_system_prompt + E_SYS + prompt + E_INST\n",
    "        \n",
    "        prompts.append(prompt)\n",
    "        answers.append(answer)\n",
    "        periods.append(period)\n",
    "        labels.append(label)\n",
    "        \n",
    "    return {\n",
    "        \"prompt\": prompts,\n",
    "        \"answer\": answers,\n",
    "        \"period\": periods,\n",
    "        \"label\": labels,\n",
    "    }\n",
    "\n",
    "def create_dataset_new(symbol_list, train_ratio=0.8, with_basics=True):\n",
    "\n",
    "    train_dataset_list = []\n",
    "    test_dataset_list = []\n",
    "\n",
    "    for symbol in symbol_list:\n",
    "\n",
    "        data_dict = gpt4_to_llama_Chinese_new(symbol, with_basics)\n",
    "#         print(data_dict['prompt'][-1])\n",
    "#         print(data_dict['answer'][-1])\n",
    "        symbols = [symbol] * len(data_dict['label'])\n",
    "        data_dict.update({\"symbol\": symbols})\n",
    "\n",
    "        dataset = Dataset.from_dict(data_dict)\n",
    "        train_size = round(train_ratio * len(dataset))\n",
    "\n",
    "        train_dataset_list.append(dataset.select(range(train_size)))\n",
    "        test_dataset_list.append(dataset.select(range(train_size, len(dataset))))\n",
    "\n",
    "    train_dataset = datasets.concatenate_datasets(train_dataset_list)\n",
    "    test_dataset = datasets.concatenate_datasets(test_dataset_list)\n",
    "\n",
    "    dataset = datasets.DatasetDict({\n",
    "        'train': train_dataset,\n",
    "        'test': test_dataset\n",
    "    })\n",
    "    \n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = create_dataset_new(tickers[:30])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a77b997963dc41fabce72aa0de140ba2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Saving the dataset (0/1 shards):   0%|          | 0/1110 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "59cf7508e5db48389dbd8ffd6eea30bc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Saving the dataset (0/1 shards):   0%|          | 0/270 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "index_name = \"SZ50\"\n",
    "min_past_weeks = 1\n",
    "max_past_weeks = 2\n",
    "train_ratio = 0.8\n",
    "dataset.save_to_disk(\n",
    "    f\"./SZdata0413/fingpt-forecaster-{index_name.lower()}-{start_date.replace('-', '')}-{end_date.replace('-', '')}-{min_past_weeks}-{max_past_weeks}-{str(train_ratio).replace('.', '')}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
